package org.znxs.znagent_s.rag;

import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;

import java.util.Collections;
import java.util.List;

/**
 * 自定义 兜底策略，采用重复检索策略，若没有检索到内容，则从备用的向量数据库中进行检索内容
 * 【注意】 不过使用这个就不能使用一些其他类指定的
 *  类似于 spring ai 官方的 VectorStoreDocumentRetriever 能够配置表达式过滤以及相似度过滤
 */
@Slf4j
public class RetryableDocumentRetriever implements DocumentRetriever {
    private final DocumentRetriever delegate; // 文档检索
    private final int maxRetries; // 最大次数
    private final VectorStore fallbackVectorStore; // 可选：备用向量库

    public RetryableDocumentRetriever(DocumentRetriever delegate, int maxRetries) {
        this.delegate = delegate;
        this.maxRetries = maxRetries;
        this.fallbackVectorStore = null;
    }

    public RetryableDocumentRetriever(DocumentRetriever delegate, int maxRetries, VectorStore fallbackVectorStore) {
        this.delegate = delegate;
        this.maxRetries = maxRetries;
        this.fallbackVectorStore = fallbackVectorStore; // 传入备用 store
    }

    @NotNull
    @Override
    public List<Document> retrieve(@NotNull Query query) {
        List<Document> results = Collections.emptyList();

        for (int i = 0; i < maxRetries; i++) {
            // 进行召回 判断召回文档数是否为0
            results = delegate.retrieve(query);
            if (!results.isEmpty()) {
                log.info("✅ 第 {} 次尝试成功召回 {} 个文档", i + 1, results.size());
                break;
            } else {
                log.warn("⚠️ 第 {} 次尝试未召回任何文档，尝试重新检索...", i + 1);
                try {
                    Thread.sleep(500); // 简单延迟
                } catch (InterruptedException ignored) {
                }
            }
        }

        // 若召回数据为空 有兜底的向量数据库
        if (results.isEmpty() && fallbackVectorStore != null) {
            log.warn("🔄 使用备用向量库进行兜底检索");
            results = List.of(new Document("#### 未成年\n" +
                    "未成年就需要玩游戏，推荐玩明日方舟"));
        }

        // 若召回数据为空 并且没有兜底备用向量数据库
        if (results.isEmpty()) {
            log.error("❎ 最终仍未召回任何文档，可返回默认内容或抛出异常");
            // 可选：抛出异常、返回默认文档等
//            return getDefaultDocuments(); // 自定义兜底文档
            throw new RuntimeException("❌ 兜底失败，请查看RAG知识库内容是否符合");
        }
        return results;
    }
}


